package cn.micro.ttl.mdc;

import com.alibaba.ttl.spi.TtlEnhanced;
import com.alibaba.ttl.threadpool.*;
import com.alibaba.ttl.threadpool.agent.TtlAgent;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.core.task.AsyncTaskExecutor;
import org.springframework.scheduling.SchedulingTaskExecutor;
import org.springframework.scheduling.TaskScheduler;

import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ScheduledExecutorService;

public class TtlExecutorsPostProcessor implements BeanPostProcessor {
    @Override
    public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
        if (bean instanceof Executor) {
            Executor executor = (Executor) bean;
            if (TtlExecutors.isTtlWrapper(executor)) {
                return executor;
            }
            if (TtlAgent.isTtlAgentLoaded() || executor instanceof TtlEnhanced) {
                return executor;
            }
            Class<?> beanClass = bean.getClass();
            if (SchedulingTaskExecutor.class.isAssignableFrom(beanClass) && TaskScheduler.class.isAssignableFrom(beanClass)) {
                return new TaskSchedulerTtlMDCWrapper((TaskScheduler) bean, (SchedulingTaskExecutor) bean, true);
            } else if (AsyncTaskExecutor.class.isAssignableFrom(beanClass)) {
                return new AsyncTaskExecutorTtlMDCWrapper((AsyncTaskExecutor) bean, true);
            } else if (ScheduledExecutorService.class.isAssignableFrom(beanClass)) {
                return new ScheduledExecutorServiceTtlMDCWrapper((ScheduledExecutorService) bean, true);
            } else if (ExecutorService.class.isAssignableFrom(beanClass)) {
                return new ExecutorServiceTtlMDCWrapper((ExecutorService) bean, true);
            } else {
                return new ExecutorTtlMDCWrapper(executor, true);
            }
        }
        return bean;
    }
}
